﻿package main

import (
	"encoding/binary"
	"io"
	"log"
	"net"

	"github.com/hashicorp/yamux"
)

type ClientBridgeHandler struct {
	TargetServer string // 本地访问接口
	BridgeServer string // tcp server 地址
	tcpRemote    *net.Conn
	tcpSession   *yamux.Session
}

func (h *ClientBridgeHandler) Login() bool {
	c, err := net.Dial("tcp", h.BridgeServer)
	if err != nil {
		return false
	}
	h.tcpRemote = &c

	// send
	cLoginReq := ClientLoginReq{
		Version: [2]byte{0x01, 0x00},
		Uid:     10010,
		Key:     [10]byte{},
	}
	myKey := []byte("1234567890")
	for i := 0; i < 10; i++ {
		cLoginReq.Key[i] = myKey[i]
	}

	log.Println("try login with %v", cLoginReq)
	if err := binary.Write(c, binary.LittleEndian, cLoginReq); err != nil {
		return false
	}

	cLoginResp := ClientLoginResp{}
	if err := binary.Read(c, binary.LittleEndian, &cLoginResp); err != nil {
		return false
	}

	log.Println("try login resp %v", cLoginResp)

	if cLoginResp.Code != 0 {
		return false
	}

	return true
}

func (h *ClientBridgeHandler) ProxyLocal() {
	defer h.tcpSession.Close()
	// yamux
	session, err := yamux.Server(*h.tcpRemote, nil)
	if err != nil {
		log.Println("yamux server error: %v", err)
		return
	}
	h.tcpSession = session
	defer session.Close()

	for {
		// new stream
		stream, err := session.Accept()
		if err != nil {
			log.Println("stream accept error: ", err)
			break
		}

		log.Println("new yamux accept to ", h.TargetServer, "streams:", session.NumStreams())

		// diag to local
		c, err := net.Dial("tcp", h.TargetServer)
		if err != nil {
			stream.Close()
			log.Println("dial to local server error: %v", err)
			continue
		}

		// copy
		go h.IoCopy(c, stream)
	}
}
func (h *ClientBridgeHandler) IoCopy(c1, c2 net.Conn) {
	defer c1.Close()
	defer c2.Close()
	log.Println("yamux stream start", c1.LocalAddr())

	// copy
	go io.Copy(c2, c1)
	io.Copy(c1, c2)

	log.Println("yamux stream close", c1.LocalAddr())
}

func main() {
	c := ClientBridgeHandler{
		TargetServer: "127.0.0.1:8080",
		BridgeServer: "127.0.0.1:9000",
	}

	if !c.Login() {
		log.Println("login error")
		return
	}

	// create stream
	c.ProxyLocal()

	log.Println("error exit")
}
